#!/usr/bin/python
#-*- coding=utf-8 -*-


from __future__ import print_function

import zmq
import json
from constant import TaskType, WorkerStatus
import time

def select_worker(socket, workers,task):
    """选择工人

    :workers: 所有工人
    :task: 任务

    """

    #TODO: 可以换成你自己想要的逻辑
    wlist = workers.get(task['task_type'], None)
    if (not wlist):
        return

    for worker in wlist:
        if (wlist[worker]['working'] >= wlist[worker]['capacity']):
            continue

        rsp = json.dumps(task)
        socket.send_multipart([worker, rsp])
        wlist[worker]['working'] += 1
        # print("proxy send to", worker, rsp)
        return

def main():
    """Load balancer main loop."""
    # Prepare context and sockets
    context = zmq.Context.instance()
    frontend = context.socket(zmq.ROUTER)
    #连接django
    frontend.bind("tcp://*:50000")
    #连接worker
    backend = context.socket(zmq.ROUTER)
    backend.bind("tcp://*:50001")

    # Initialize main loop state
    workers = {}
    tasks = {}
    task_idx = 1
    poller = zmq.Poller()
    # Only poll for requests from backend until workers are available
    poller.register(backend, zmq.POLLIN)
    poller.register(frontend, zmq.POLLIN)

    while True:
        sockets = dict(poller.poll())

        if backend in sockets:
            # Handle worker activity on the backend
            worker, req = backend.recv_multipart()
            req = json.loads(req)
            # print("proxy recv from", worker, req)

            wlist = workers.get(req['task_type'], {})
            wlist[worker] = req
            workers[req['task_type']] = wlist

            if (req.get('task', None) and req.get('idx', None)):
                del tasks[req['idx']]

        if frontend in sockets:
            # Get next client request, route to last-used worker
            client, req = frontend.recv_multipart()
            # print("proxy recv from", client, req)

            req = json.loads(req)
            req['client'] = client
            req['idx'] = task_idx
            tasks[task_idx] = req
            task_idx = task_idx + 1

        for idx in tasks:
            select_worker(backend, workers, tasks[idx])

    # Clean up
    backend.close()
    frontend.close()
    context.term()


if __name__ == "__main__":
    main()
